package cn.android666.mediacodecsynchronous

import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import timber.log.Timber


@Suppress("PrivatePropertyName")
class H264Encoder(videoWidth: Int, videoHeight: Int) {

    private val mMediaCodec: MediaCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC)
    private var mBufferInfo = MediaCodec.BufferInfo()
    private val _1K = 1000
    private val _1M = _1K * 1000
    private val ySize = videoWidth * videoHeight // Y分量大小
    private val oneFrameSize = (ySize * 3) shr 1 // 一帧画面大小
    private var index: Int = 0
    private var temp: Byte = 0
    private val h264Saver: H264Saver by lazy { H264Saver() }
    private var isPutEmptyArray = false

    init {
        val mediaFormat = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, videoWidth, videoHeight)
        mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible)
        mediaFormat.setInteger(MediaFormat.KEY_BITRATE_MODE, MediaCodecInfo.EncoderCapabilities.BITRATE_MODE_VBR) // 设置码率为动态码率，默认也是动态码率
        mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, getBitrate(videoWidth, videoHeight)) // 码率（即比特率）, 官方Demo这里的1000即1kbps，1000_000即1mpbs
        mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, 25)      // 帧速（25帧/秒）
        mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 2) // I帧间隔（1帧/2秒），因为帧 1秒出25帧，2秒就出50帧，所以I帧间隔为2的话就是每50帧出一个关键帧

        // 第二个参数用于显示解码器的视频内容，第三个参数为编解码器的解密参数，第四个参数为指定为编码器
        mMediaCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
        mMediaCodec.start()
        Timber.i("实际使用的H264编码器：${mMediaCodec.codecInfo.name}")
    }

    /** 把NV21格式的Yuv数据编码为H264数据 */
    fun encodeYuvToH264(yuvBytes: ByteArray) {
        if (isPutEmptyArray) return // 已经放入了空数组（用于结束编码），则不处理，因为多线程，所以有可能释放的时候还有数据扔进来编码的

        val flags = if (yuvBytes.isEmpty()) {
            isPutEmptyArray = true
            MediaCodec.BUFFER_FLAG_END_OF_STREAM
        } else {
            nv21ToNv12(yuvBytes)
            0
        }

        val inputBufferIndex = mMediaCodec.dequeueInputBuffer(10_000) // 如果10毫秒都等不到可用缓冲，则这一帧的yuv数据将丢掉。谷歌官方Demo也是用的这个值
        if (inputBufferIndex >= 0) {
            val inputBuffer = mMediaCodec.getInputBuffer(inputBufferIndex) ?: return
            inputBuffer.put(yuvBytes, 0, yuvBytes.size) // 官方Demo在调用put之前会先调用inputBuffer.clear()，实际上并不需要
            mMediaCodec.queueInputBuffer(inputBufferIndex,0, yuvBytes.size,System.nanoTime() / 1000, flags)
        }

        // 从MediaCodec中取出编好的H264数据并使用（如保存、发送）
        var outputBufferIndex: Int
        while (mMediaCodec.dequeueOutputBuffer(mBufferInfo, 10_000).also { outputBufferIndex = it } >= 0) {
            val outputBuffer = mMediaCodec.getOutputBuffer(outputBufferIndex) ?: return
            when (mBufferInfo.flags) {
                MediaCodec.BUFFER_FLAG_CODEC_CONFIG, MediaCodec.BUFFER_FLAG_KEY_FRAME, 0 -> { // 配置帧、关键帧、普通帧
                    h264Saver.write(outputBuffer)
                    mMediaCodec.releaseOutputBuffer(outputBufferIndex, false)
                }
                MediaCodec.BUFFER_FLAG_END_OF_STREAM -> {
                    Timber.i("已经到达流的终点了")
                    mMediaCodec.releaseOutputBuffer(outputBufferIndex, false)
                    releaseMediaCodec()
                    break // 退出循环，无需再去获取编码的数据。
                }
            }
        }
    }

    private fun getBitrate(videoWidth: Int, videoHeight: Int): Int = when {
        (videoWidth == 1920 && videoHeight == 1080) || (videoWidth == 1080 && videoHeight == 1920) -> _1M * 3 shr 1
        (videoWidth == 1280 && videoHeight == 720) || (videoWidth == 720 && videoHeight == 1280) -> _1M
        (videoWidth == 640 && videoHeight == 480) || (videoWidth == 480 && videoHeight == 640) -> _1K * 500
        (videoWidth == 352 && videoHeight == 288) || (videoWidth == 288 && videoHeight == 352) -> _1K * 300
        else -> _1M * 1
    }

    private fun nv21ToNv12(yuvBytes: ByteArray) {
        index = ySize
        while (index < oneFrameSize) {
            temp = yuvBytes[index]
            yuvBytes[index] = yuvBytes[index + 1]
            yuvBytes[index + 1] = temp
            index += 2
        }
    }

    /** 关闭编码器 */
    fun close() {
        Timber.i("close")
        encodeYuvToH264(ByteArray(0))
    }

    private fun releaseMediaCodec() {
        try {
            mMediaCodec.stop()
        } catch (e: Exception) {
            Timber.e(e,"别慌，正常停止${H264Encoder::class.java.simpleName}时出现的异常！")
        }

        try {
            mMediaCodec.release()
        } catch (e: Exception) {
            Timber.e(e,"别慌，正常释放${H264Encoder::class.java.simpleName}时出现的异常！")
        }

        h264Saver.close()
    }

}
